import os
import subprocess
from fastapi import FastAPI, HTTPException, File, UploadFile
from pathlib import Path
import zipfile
import shutil

app = FastAPI()

# 路径设置
MODEL_DIR = Path("models")
OUTPUT_DIR = Path("converted_models")

# 确保目录存在
MODEL_DIR.mkdir(parents=True, exist_ok=True)
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)


@app.post("/convert_model/")
async def convert_model(input_model: UploadFile = File(...), output_format: str = 'keras'):
    """
    接口：接收一个上传的 TensorFlow.js 模型的压缩包（.zip 格式），
    解压并转换为指定格式的 Keras 模型。

    - **input_model**: 上传的 TensorFlow.js 模型压缩包（.zip 格式）。
    - **output_format**: 转换后的格式（默认是 keras）。
    """

    # 保存上传的 zip 文件
    zip_path = MODEL_DIR / input_model.filename
    with open(zip_path, 'wb') as f:
        content = await input_model.read()
        f.write(content)

    # 解压缩 zip 文件
    unzip_dir = MODEL_DIR / input_model.filename.split('.')[0]
    if unzip_dir.exists():
        shutil.rmtree(unzip_dir)  # 如果目录已经存在，先删除
    os.makedirs(unzip_dir)

    # 解压 .zip 文件
    try:
        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
            zip_ref.extractall(unzip_dir)
    except zipfile.BadZipFile:
        raise HTTPException(status_code=400, detail="Uploaded file is not a valid ZIP archive.")

    # 确保解压出来的文件包含 model.json
    model_json_path = unzip_dir / 'model.json'
    if not model_json_path.exists():
        raise HTTPException(status_code=400, detail="model.json not found in the uploaded ZIP file.")

    # 设置输出文件路径为 .h5 文件
    output_model_path = OUTPUT_DIR / f"{input_model.filename.split('.')[0]}.h5"

    # 设置转换命令
    command = [
        "tensorflowjs_converter",
        "--input_format=tfjs_layers_model",
        f"--output_format={output_format}",
        str(model_json_path),
        str(output_model_path)  # 输出为 .h5 文件
    ]

    try:
        # 调用 tensorflowjs_converter 进行转换
        subprocess.run(command, check=True)

        # 返回转换后的模型路径
        return {"message": f"Model converted successfully!", "converted_model_path": str(output_model_path)}

    except subprocess.CalledProcessError as e:
        raise HTTPException(status_code=400, detail=f"Model conversion failed: {e}")
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/")
async def root():
    return {"message": "Welcome to the TensorFlow.js Model Converter API!"}


from fastapi.middleware.cors import CORSMiddleware
# 允许跨域访问
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 设置为 "*" 允许所有来源访问
    allow_methods=["*"],
    allow_headers=["*"],
)

if __name__ == '__main__':
    import uvicorn

    uvicorn.run(app, host='0.0.0.0', port=18888)